#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
将SQLite数据库中的表迁移到MySQL数据库
"""

import os
import sys
import sqlite3
import pymysql
import argparse
from tqdm import tqdm

def get_table_structure(sqlite_cursor, table_name):
    """获取SQLite表结构"""
    # 获取表的列信息
    sqlite_cursor.execute(f"PRAGMA table_info({table_name})")
    columns = sqlite_cursor.fetchall()
    return [column[1] for column in columns]  # 返回列名列表

def get_table_data(sqlite_cursor, table_name):
    """获取SQLite表中的所有数据"""
    sqlite_cursor.execute(f"SELECT * FROM {table_name}")
    return sqlite_cursor.fetchall()

def create_insert_query(table_name, columns, values):
    """创建MySQL的插入语句"""
    placeholders = ', '.join(['%s'] * len(values))
    # 给所有列名添加反引号，防止与MySQL保留字冲突
    columns_str = ', '.join([f'`{column}`' for column in columns])
    query = f"INSERT INTO `{table_name}` ({columns_str}) VALUES ({placeholders})"
    return query

def process_data_for_mysql(row):
    """处理数据使其适合MySQL插入
    
    - 将None值替换为适合MySQL的NULL
    - 处理布尔值
    - 其他特殊类型转换
    """
    processed_row = []
    for value in row:
        # 处理None值
        if value is None:
            processed_row.append(None)
        # 处理布尔值
        elif isinstance(value, bool):
            processed_row.append(1 if value else 0)
        # 其他值保持不变
        else:
            processed_row.append(value)
    return tuple(processed_row)

def migrate_table(sqlite_db_path, table_name, mysql_config):
    """将SQLite表迁移到MySQL"""
    try:
        # 连接SQLite数据库
        sqlite_conn = sqlite3.connect(sqlite_db_path)
        sqlite_cursor = sqlite_conn.cursor()
        
        # 检查表是否存在
        sqlite_cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name=?", (table_name,))
        if not sqlite_cursor.fetchone():
            print(f"错误: 表 '{table_name}' 在SQLite数据库中不存在")
            return False
        
        # 获取表结构和数据
        columns = get_table_structure(sqlite_cursor, table_name)
        data = get_table_data(sqlite_cursor, table_name)
        
        if not data:
            print(f"注意: 表 '{table_name}' 没有数据可以迁移")
            return True
        
        # 连接MySQL数据库
        mysql_conn = pymysql.connect(
            host=mysql_config['host'],
            user=mysql_config['user'],
            password=mysql_config['password'],
            database=mysql_config['database'],
            charset='utf8mb4'
        )
        mysql_cursor = mysql_conn.cursor()
        
        # 创建插入语句
        insert_query = create_insert_query(table_name, columns, data[0])
        print(f"SQL语句: {insert_query}")  # 调试信息
        
        # 批量插入数据
        print(f"开始迁移表 '{table_name}' 的 {len(data)} 条记录...")
        
        # 使用tqdm显示进度条
        success_count = 0
        error_count = 0
        for row in tqdm(data, desc=f"迁移 {table_name}"):
            try:
                # 处理数据使其适合MySQL
                processed_row = process_data_for_mysql(row)
                mysql_cursor.execute(insert_query, processed_row)
                success_count += 1
            except pymysql.err.IntegrityError as e:
                # 如果有主键冲突，打印警告但继续执行
                print(f"警告: 记录插入失败 - {e}")
                error_count += 1
                continue
            except Exception as e:
                print(f"错误: 记录插入失败 - {e}")
                print(f"问题数据: {row}")
                error_count += 1
                if error_count > 5:
                    raise Exception(f"连续错误次数过多，终止迁移")
        
        # 提交更改
        mysql_conn.commit()
        
        # 关闭连接
        mysql_cursor.close()
        mysql_conn.close()
        sqlite_cursor.close()
        sqlite_conn.close()
        
        print(f"表 '{table_name}' 迁移完成! 成功: {success_count}, 失败: {error_count}")
        return True
        
    except Exception as e:
        print(f"迁移过程中出错: {e}")
        return False

def main():
    parser = argparse.ArgumentParser(description='将SQLite数据库表迁移到MySQL')
    parser.add_argument('--sqlite', required=True, help='SQLite数据库文件路径')
    parser.add_argument('--table', required=True, help='要迁移的表名')
    parser.add_argument('--mysql-host', default='localhost', help='MySQL主机地址')
    parser.add_argument('--mysql-user', required=True, help='MySQL用户名')
    parser.add_argument('--mysql-password', required=True, help='MySQL密码')
    parser.add_argument('--mysql-db', required=True, help='MySQL数据库名')
    
    args = parser.parse_args()
    
    # 检查SQLite文件是否存在
    if not os.path.exists(args.sqlite):
        print(f"错误: SQLite数据库文件 '{args.sqlite}' 不存在")
        return 1
    
    mysql_config = {
        'host': args.mysql_host,
        'user': args.mysql_user,
        'password': args.mysql_password,
        'database': args.mysql_db
    }
    
    success = migrate_table(args.sqlite, args.table, mysql_config)
    return 0 if success else 1

if __name__ == "__main__":
    sys.exit(main()) 